use pnet::datalink::{self, DataLinkSender};
use pnet::datalink::Channel::Ethernet;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket, MutableEthernetPacket};
use pnet::packet::ip::IpNextHeaderProtocols;
use pnet::packet::ipv4::{Ipv4, Ipv4Packet, MutableIpv4Packet};
use pnet::packet::tcp::{ipv4_checksum, MutableTcpPacket, Tcp, TcpPacket};
use pnet::packet::Packet;

fn main() {
    // 获取网络接口
    let interfaces = datalink::interfaces();
    let interface = match interfaces.iter()
        .find(|iface| iface.is_up() && !iface.is_loopback() && !iface.ips.is_empty())
    {
        Some(iface) => iface,
        None => {
            eprintln!("No suitable network interface found.");
            return;
        }
    };

    println!("Monitoring interface: {}", interface.name);

    // 创建数据链路层通道
    let (mut tx, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("Unsupported channel type"),
        Err(e) => panic!("Failed to create channel: {}", e),
    };

    loop {
        match rx.next() {
            Ok(packet) => {
                if let Some(eth_packet) = EthernetPacket::new(packet) {
                    match eth_packet.get_ethertype() {
                        EtherTypes::Ipv4 => {
                            if let Some(ipv4_packet) = Ipv4Packet::new(eth_packet.payload()) {
                                if let Some(modified) = handle_ipv4_packet(&ipv4_packet) {
                                    send_modified_packet(tx.as_mut(), &eth_packet, &modified);
                                }
                            }
                        }
                        _ => (),
                    }
                }
            }
            Err(e) => eprintln!("Packet receive error: {}", e),
        }
    }
}

fn handle_ipv4_packet(ipv4_packet: &Ipv4Packet) -> Option<Vec<u8>> {
    let mut buffer = vec![0; ipv4_packet.packet().len()];
    let mut mut_ipv4 = MutableIpv4Packet::new(&mut buffer)?;

    let p = ipv4_packet.packet();
    // 复制原始 IPv4 头部
    mut_ipv4.populate(&Ipv4 {
        version: ipv4_packet.get_version(),
        header_length: ipv4_packet.get_header_length(),
        dscp: ipv4_packet.get_dscp(),
        ecn: ipv4_packet.get_ecn(),
        total_length: ipv4_packet.get_total_length(),
        identification: ipv4_packet.get_identification(),
        flags: ipv4_packet.get_flags(),
        fragment_offset: ipv4_packet.get_fragment_offset(),
        next_level_protocol: ipv4_packet.get_next_level_protocol(),
        source: ipv4_packet.get_source(),
        destination: ipv4_packet.get_destination(),
        options: ipv4_packet.get_options().to_vec(),
        checksum: 0,
        ttl: ipv4_packet.get_ttl(),
        payload: ipv4_packet.payload().to_vec(),
    });

    match ipv4_packet.get_next_level_protocol() {
        IpNextHeaderProtocols::Tcp => {
            if let Some(tcp_packet) = TcpPacket::new(ipv4_packet.payload()) {
                let mut tcp_buffer = vec![0; tcp_packet.packet().len()];
                let mut mut_tcp = MutableTcpPacket::new(&mut tcp_buffer)?;

                // 修改 payload（示例：前 8 字节改为 'X'）
                let mut payload = tcp_packet.payload().to_vec();
                if payload.len() >= 8 {
                    payload[..8].copy_from_slice(b"XXXXXXXX");
                }
                // 复制原始 TCP 头部
                mut_tcp.populate(&Tcp {
                    source: tcp_packet.get_source(),
                    destination: tcp_packet.get_destination(),
                    sequence: tcp_packet.get_sequence(),
                    acknowledgement: tcp_packet.get_acknowledgement(),
                    data_offset: tcp_packet.get_data_offset(),
                    reserved: tcp_packet.get_reserved(),
                    flags: tcp_packet.get_flags(),
                    window: tcp_packet.get_window(),
                    checksum: 0,
                    options: tcp_packet.get_options().to_vec(),
                    payload: payload,
                    urgent_ptr: tcp_packet.get_urgent_ptr(),
                });


                // 重新计算 TCP 校验和
                let checksum = ipv4_checksum(
                    &mut_tcp.to_immutable(),
                    &ipv4_packet.get_source(),
                    &ipv4_packet.get_destination(),
                );
                mut_tcp.set_checksum(checksum);

                // 更新 IPv4 数据包
                mut_ipv4.set_payload(mut_tcp.packet());
                mut_ipv4.set_checksum(0); // 强制重新计算
            }
        }
        _ => return None,
    }

    Some(buffer)
}

fn send_modified_packet(tx: &mut dyn DataLinkSender, original_eth: &EthernetPacket, modified_ip_payload: &[u8]) {
    let original_len = original_eth.packet().len();
    let mut eth_buffer = vec![0; original_eth.packet().len()];
    let mut eth_packet = MutableEthernetPacket::new(&mut eth_buffer).unwrap();

    // 重建以太网帧
    eth_packet.set_destination(original_eth.get_destination());
    eth_packet.set_source(original_eth.get_source());
    eth_packet.set_ethertype(original_eth.get_ethertype());
    eth_packet.set_payload(modified_ip_payload);

    let len = eth_packet.packet().len();
    println!("original_len: {}, len: {}", original_len, len);

    if let Some(Err(e)) = tx.send_to(eth_packet.packet(), None) {
        eprintln!("Failed to send packet: {}", e);
    }
}